package video

import (
	"encoding/json"
	"errors"
	"fmt"
	"geekai/core/store"
	"geekai/core/types"
	logger2 "geekai/logger"
	"geekai/model"
	"geekai/service"
	"geekai/service/oss"
	"geekai/utils"
	"github.com/go-redis/redis/v8"
	"github.com/imroc/req/v3"
	"gorm.io/gorm"
	"io"
	"time"
)

const (
	// LumaGenerationsGetTask Luma 查询单个任务
	LumaGenerationsGetTask = "%s/luma/generations/%s"
	// LumaGenerationsCreate Luma 生成视频
	LumaGenerationsCreate = "%s/luma/generations"
)

var log = logger2.GetLogger()

type Service struct {
	httpClient  *req.Client
	taskQueue   *store.RedisQueue
	notifyQueue *store.RedisQueue
	db          *gorm.DB
	uploader    *oss.UploaderManager
	wsService   *service.WebSocketService
	userService *service.UserService
	clientIds   map[uint]string
}

func NewService(db *gorm.DB, manager *oss.UploaderManager,
	redisCli *redis.Client, wsService *service.WebSocketService,
	userService *service.UserService) *Service {
	return &Service{
		httpClient:  req.C().SetTimeout(time.Minute * 3),
		taskQueue:   store.NewRedisQueue("My_Luma_Task_Queue", redisCli),
		notifyQueue: store.NewRedisQueue("My_Luma_Notify_Queue", redisCli),
		db:          db,
		wsService:   wsService,
		uploader:    manager,
		userService: userService,
		clientIds:   map[uint]string{},
	}
}

type LumaRespVo struct {
	Id                  string      `json:"id"`
	Prompt              string      `json:"prompt"`
	State               string      `json:"state"`
	QueueState          interface{} `json:"queue_state"`
	CreatedAt           string      `json:"created_at"`
	Video               interface{} `json:"video"`
	VideoRaw            interface{} `json:"video_raw"`
	Liked               interface{} `json:"liked"`
	EstimateWaitSeconds interface{} `json:"estimate_wait_seconds"`
	Thumbnail           interface{} `json:"thumbnail"`
	Channel             string      `json:"channel,omitempty"`
}

type LumaTaskVo struct {
	Id    string      `json:"id"`
	Liked interface{} `json:"liked"`
	State string      `json:"state"`
	Video struct {
		Url         string `json:"url"`
		Width       int    `json:"width"`
		Height      int    `json:"height"`
		Thumbnail   string `json:"thumbnail"`
		DownloadUrl string `json:"download_url"`
	} `json:"video"`
	Prompt    string `json:"prompt"`
	UserId    string `json:"user_id"`
	BatchId   string `json:"batch_id"`
	Thumbnail struct {
		Url    string `json:"url"`
		Width  int    `json:"width"`
		Height int    `json:"height"`
	} `json:"thumbnail"`
	VideoRaw struct {
		Url    string `json:"url"`
		Width  int    `json:"width"`
		Height int    `json:"height"`
	} `json:"video_raw"`
	CreatedAt string `json:"created_at"`
	LastFrame struct {
		Url    string `json:"url"`
		Width  int    `json:"width"`
		Height int    `json:"height"`
	} `json:"last_frame"`
}

func (s *Service) Run() {
	// 将数据库中未提交的任务加载到队列中
	var jobs []model.VideoJob
	s.db.Where("task_id", "").Where("progress", 0).Find(&jobs)
	for _, job := range jobs {
		var task types.VideoTask
		err := utils.JsonDecode(job.TaskInfo, &task)
		if err != nil {
			log.Errorf("Luma 解码任务信息时出错：%v", err)
			continue
		}
		task.Id = job.Id
		s.PushTask(task)
		s.clientIds[job.Id] = task.ClientId
	}
	log.Info("启动Luma任务消费者 ...")
	go func() {
		for {
			var task types.VideoTask
			err := s.taskQueue.LPop(&task)
			if err != nil {
				log.Errorf("Luma 消息任务队列任务出现错误：%v", err)
				continue
			}

			// 提示词翻译
			if utils.HasChinese(task.Prompt) {
				content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), task.TranslateModelId)
				if err == nil {
					task.Prompt = content
				} else {
					log.Warnf("Luma 提示词进行翻译出错：%v", err)
				}
			}

			if task.ClientId != "" {
				s.clientIds[task.Id] = task.ClientId
			}
			var r LumaRespVo
			r, err = s.LumaCreate(task)
			if err != nil {
				log.Errorf("Luma 生成视频任务执行失败：%v", err)
				err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
					"err_msg":   err.Error(),
					"progress":  service.FailTaskProgress,
					"cover_url": "/images/failed.jpg",
				}).Error
				if err != nil {
					log.Errorf("Luma 更新数据库任务信息失败：%v", err)
				}
				s.notifyQueue.RPush(service.NotifyMessage{
					UserId:   task.UserId,
					ClientId: task.ClientId,
					JobId:    int(task.Id),
					Message:  service.TaskStatusFailed,
				})
				continue
			}
			// 更新任务信息
			err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
				"task_id":    r.Id,
				"channel":    r.Channel,
				"prompt_ext": r.Prompt,
			}).Error
			if err != nil {
				log.Errorf("Luma 更新数据库任务信息失败：%v", err)
				s.PushTask(task)
			}
		}
	}()
}

// LumaCreate 生成视频
func (s *Service) LumaCreate(task types.VideoTask) (LumaRespVo, error) {
	// 读取 API KEY
	var apiKey model.ApiKey
	session := s.db.Session(&gorm.Session{}).Where("type", "luma").Where("enabled", true)
	if task.Channel != "" {
		session = session.Where("api_url", task.Channel)
	}
	tx := session.Order("last_used_at desc").First(&apiKey)
	if tx.Error != nil {
		return LumaRespVo{}, errors.New("没有可用于Suno的API密钥")
	}
	reqBody := map[string]interface{}{
		"user_prompt":   task.Prompt,
		"expand_prompt": task.Params.PromptOptimize, // 是否优化提示词
		"loop":          task.Params.Loop,           // 是否循环参考图
		"image_url":     task.Params.StartImgURL,    // 第一帧参考图地址
		"image_end_url": task.Params.EndImgURL,      // 最后一帧参考图地址
	}
	var res LumaRespVo
	apiURL := fmt.Sprintf(LumaGenerationsCreate, apiKey.ApiURL)
	log.Debugf("Luma：请求地址：%s，请求参数：%+v", apiURL, reqBody)
	r, err := s.httpClient.R().
		SetHeader("Authorization", "Bearer "+apiKey.Value).
		SetBody(reqBody).
		Post(apiURL)
	if err != nil {
		return LumaRespVo{}, fmt.Errorf("Luma请求API：%s 出错：%v", apiURL, err)
	}

	if r.StatusCode != 200 && r.StatusCode != 201 {
		return LumaRespVo{}, fmt.Errorf("Luma请求API：%s 出错：%v", apiURL, r.String())
	}

	body, _ := io.ReadAll(r.Body)
	err = json.Unmarshal(body, &res)
	if err != nil {
		return LumaRespVo{}, fmt.Errorf("Luma解析API返回数据失败：%v, %s", err, string(body))
	}

	// 更新 API KEY 最后一次使用时间
	apiKey.LastUsedAt = time.Now().Unix()
	session.Updates(&apiKey)
	// 设置通道为，我们这次请求使用的API KEY 对应的 API 地址
	res.Channel = apiKey.ApiURL
	return res, nil
}

// CheckTaskNotify 监听消息通知队列
func (s *Service) CheckTaskNotify() {
	go func() {
		log.Info("正在运行Luma任务通知检查 ...")
		for {
			var message service.NotifyMessage
			err := s.notifyQueue.LPop(&message)
			if err != nil {
				log.Errorf("Luma 从消息通知队列消费消息出现错误：%v", err)
				continue
			}
			log.Debugf("Luma 消息通知：%+v", message)
			// 获取客户端连接
			client := s.wsService.Clients.Get(message.ClientId)
			if client == nil {
				log.Warnf("Luma 客户端连接为空：%+v", s.wsService.Clients)
				continue
			}
			// 发送消息
			utils.SendChannelMsg(client, types.ChLuma, message.Message)
		}
	}()
}

// DownloadFiles 下载文件
func (s *Service) DownloadFiles() {
	go func() {
		var items []model.VideoJob
		for {
			// 102 标识任务执行完成，但是文件未下载
			res := s.db.Where("progress", 102).Find(&items)
			if res.Error != nil {
				log.Errorf("Luma 下载文件时，查询数据库任务出错：%v", res.Error)
				time.Sleep(time.Second * 5)
				continue
			}
			for _, v := range items {
				if v.WaterURL == "" {
					continue
				}
				// 下载有水印视频
				log.Infof("Luma 尝试下载有水印视频：%s", v.WaterURL)
				videoURL, err := s.uploader.GetUploadHandler().PutUrlFile(v.WaterURL, true)
				if err != nil {
					log.Errorf("Luma 尝试下载有水印视频失败：%v", err)
					continue
				}
				log.Infof("Luma 下载有水印视频文件成功：%s", videoURL)
				v.WaterURL = videoURL
				// 下载无水印视频
				if v.VideoURL != "" {
					log.Infof("Luma 尝试下载无水印视频：%s", v.VideoURL)
					videoURL, err = s.uploader.GetUploadHandler().PutUrlFile(v.VideoURL, true)
					if err != nil {
						log.Errorf("Luma 尝试下载无水印视频失败：%v", err)
						continue
					}
				}
				log.Infof("Luma 下载无水印视频文件成功：%s", videoURL)
				// 更新任务信息
				v.VideoURL = videoURL
				v.Progress = 100
				s.db.Updates(&v)
				// 通知前端
				s.notifyQueue.RPush(service.NotifyMessage{
					ClientId: s.clientIds[v.Id],
					UserId:   int(v.UserId),
					JobId:    int(v.Id),
					Message:  service.TaskStatusFinished,
				})
			}
			time.Sleep(time.Second * 10)
		}
	}()
}

// SyncTaskProgress 异步拉取任务，判断任务是否执行成功
func (s *Service) SyncTaskProgress() {
	go func() {
		var jobs []model.VideoJob
		for {
			// 查询进度小于100，并且任务Id不为空的任务
			res := s.db.Where("progress < ?", 100).Where("task_id <> ?", "").Find(&jobs)
			if res.Error != nil {
				log.Errorf("Luma SyncTaskProgress 查询数据库任务失败：%v", res.Error.Error())
				continue
			}
			for _, job := range jobs {
				// 查询任务信息
				task, err := s.QueryLumaTask(job.TaskId, job.Channel)
				if err != nil {
					log.Errorf("Luma 查询任务失败：%v", err.Error())
					// 更新任务信息
					s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
						"progress": service.FailTaskProgress,
						"err_msg":  err.Error(),
					})
					continue
				}

				log.Debugf("Luma 任务返回数据：%+v", task)
				// 任务完成，删除旧的任务插入两条新任务
				if task.State == "completed" { // 更新任务信息
					data := map[string]interface{}{
						"progress":   102, // 102表示资源未下载
						"water_url":  task.Video.Url,
						"raw_data":   utils.JsonEncode(task),
						"prompt_ext": task.Prompt,
						"cover_url":  task.Thumbnail.Url,
					}
					if task.Video.DownloadUrl != "" {
						data["video_url"] = task.Video.DownloadUrl
					}
					err = s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(data).Error
					if err != nil {
						log.Errorf("Luma 更新数据库任务失败：%v", err)
						continue
					}
				}
			}
			// 找出失败的任务，并恢复其扣减的算力
			s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
			for _, job := range jobs {
				err := s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
					Type:   types.PowerRefund,
					Model:  "luma",
					Remark: fmt.Sprintf("Luma 任务失败，退回算力。任务Id：%s，Err：%s", job.TaskId, job.ErrMsg),
				})
				if err != nil {
					log.Errorf("Luma 任务失败，退回算力时出现错误：%v", err.Error())
					continue
				}
				// 更新任务状态
				s.db.Model(&job).UpdateColumn("power", 0)
			}
			time.Sleep(time.Second * 10)
		}
	}()
}

func (s *Service) PushTask(task types.VideoTask) {
	log.Infof("将新的 Luma 任务到添加到任务列表：%+v", task)
	s.taskQueue.RPush(task)
}

// QueryLumaTask 查询单个Luma任务
func (s *Service) QueryLumaTask(taskId string, channel string) (LumaTaskVo, error) {
	var apiKey model.ApiKey
	err := s.db.Session(&gorm.Session{}).Where("type", "luma").
		Where("api_url", channel).
		Where("enabled", true).
		Order("last_used_at desc").First(&apiKey).Error
	if err != nil {
		return LumaTaskVo{}, errors.New("没有可用于Luma的API密钥")
	}
	apiURL := fmt.Sprintf(LumaGenerationsGetTask, apiKey.ApiURL, taskId)
	var res LumaTaskVo
	// 发送请求
	r, err := s.httpClient.R().SetHeader("Authorization", "Bearer "+apiKey.Value).Get(apiURL)
	if err != nil {
		return LumaTaskVo{}, fmt.Errorf("Luma请求API：%s 出错：%v", apiURL, err)
	}
	defer r.Body.Close()

	if r.StatusCode != 200 {
		return LumaTaskVo{}, fmt.Errorf("Luma请求API：%s 返回失败：%v", apiURL, r.String())
	}

	body, _ := io.ReadAll(r.Body)
	err = json.Unmarshal(body, &res)
	if err != nil {
		return LumaTaskVo{}, fmt.Errorf("Luma解析API返回数据失败：%v, %s", err, string(body))
	}
	return res, nil
}
